/*
This code takes observed reports from a variable, and minimises the MSE of a KS model. 

takes data in. 
uses penalty weight on theta. 
*/


mata        //invoking MATA
mata clear  //clear MATA without disturbing Stata

void ks_function(todo,p,y,g,negH)
{
    eta = p[1]
    sigma = p[2]
    ro = p[3]
    
    // Get observed data from Stata
    observed_data = st_data(., "reports")
    
    // XL[ro] = Ceiling[ro]
    XL_ro = ceil(ro)
    
    // Conditional logic based on x >= XL[ro]
    x5_hat = 5 >= XL_ro ? eta/ ((6) * (eta - 5 + ro)) : 1/(6) *(1 - (2*normal((ro - 5)/sigma) - 1))
    x4_hat = 4 >= XL_ro ? eta/ ((6) * (eta - 4 + ro)) : 1/(6) *(1 - (2*normal((ro - 4)/sigma) - 1))
    x3_hat = 3 >= XL_ro ? eta/ ((6) * (eta - 3 + ro)) : 1/(6) *(1 - (2*normal((ro - 3)/sigma) - 1))
    x2_hat = 2 >= XL_ro ? eta/ ((6) * (eta - 2 + ro)) : 1/(6) *(1 - (2*normal((ro - 2)/sigma) - 1))
    x1_hat = 1 >= XL_ro ? eta/ ((6) * (eta - 1 + ro)) : 1/(6) *(1 - (2*normal((ro - 1)/sigma) - 1))
    x0_hat = 0 >= XL_ro ? eta/ ((6) * (eta - 0 + ro)) : 1/(6) *(1 - (2*normal((ro - 0)/sigma) - 1))
    
    // Theta function
    sum1 = 0
    for (y_val = 0; y_val <= XL_ro - 1; y_val++) {
        sum1 = sum1 + ((ro - y_val) <= 0 ? 0 : (2*normal((ro - y_val)/sigma) - 1))
    }
    
    sum2 = 0
    for (x_val = XL_ro; x_val <= 5; x_val++) {
        sum2 = sum2 + (x_val - ro) / (eta - x_val + ro)
    }
    
    theta = sum1 - sum2
    
    //** MSE calculation
	mse = mean(((observed_data[1..6] - (x0_hat \ x1_hat \ x2_hat \ x3_hat \ x4_hat \ x5_hat)) * 100):^2)
// older version     mse = ((observed_data[1]-x0_hat)^2+(observed_data[2]-x1_hat)^2+(observed_data[3]-x2_hat)^2+(observed_data[4]-x3_hat)^2+(observed_data[5]-x4_hat)^2+(observed_data[6]-x5_hat)^2)/6
    
    //** Add penalty for theta constraint (theta = 0)
    penalty_weight = 100000
    y = mse^2 + penalty_weight * theta^2
	
}

S1=optimize_init()
optimize_init_technique(S1, "nr dfp bfgs")
optimize_init_which(S1,"min")  //minimization instead of maximization
optimize_init_evaluator(S1,&ks_function())
optimize_init_params(S1,(3.65,2.96,3.05))
x=optimize(S1)
x
p = optimize_result_value(S1)
p   // mse
x[1] // eta
x[2] // sigma 
x[3] // ro
st_numscalar("eta", x[1])
st_numscalar("sigma", x[2])
st_numscalar("ro", x[3])

// Calculate final theta with optimal parameters
eta_opt = x[1]; sigma_opt = x[2]; ro_opt = x[3]
XL_ro_opt = ceil(ro_opt)
sum1_opt = 0; for(y=0; y<=XL_ro_opt-1; y++) sum1_opt = sum1_opt + ((ro_opt-y)<=0 ? 0 : (2*normal((ro_opt-y)/sigma_opt)-1))
sum2_opt = 0; for(x_val=XL_ro_opt; x_val<=5; x_val++) sum2_opt = sum2_opt + (x_val-ro_opt)/(eta_opt-x_val+ro_opt)
theta_final = sum1_opt - sum2_opt
printf("Final theta = %9.6f\n", theta_final)

// Calculate final x*_hat values with optimal parameters
x5_hat_final = 5 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 5 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 5)/sigma_opt) - 1))
x4_hat_final = 4 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 4 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 4)/sigma_opt) - 1))
x3_hat_final = 3 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 3 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 3)/sigma_opt) - 1))
x2_hat_final = 2 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 2 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 2)/sigma_opt) - 1))
x1_hat_final = 1 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 1 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 1)/sigma_opt) - 1))
x0_hat_final = 0 >= XL_ro_opt ? eta_opt/ ((6) * (eta_opt - 0 + ro_opt)) : 1/(6) *(1 - (2*normal((ro_opt - 0)/sigma_opt) - 1))

// Store predictions as Stata variable
predictions_ks = (x0_hat_final \ x1_hat_final \ x2_hat_final \ x3_hat_final \ x4_hat_final \ x5_hat_final) * 100
st_addvar("double", "predictions_ks")
st_store(., "predictions_ks", predictions_ks)

// mata describe //show MATA enviornment
end



